#include <src/pcps.h>
#include <src/input.h>
#include <src/hamiltonian.h>
#include <src/wavefunction.h>
#include <src/cluster.h>
#include <src/coeff_value.h>
#include <src/blas_lapack.h>
#include <src/properties.h>
#include <src/mc_energy.h>
#include <src/time_evolution.h>
#include <src/timing.h>

//--------------------------------------------------------------------------------
// pcps::make_time_evolution_op -- Adds the contribution of a hamiltonian operator
//                                 to one or more time evolution operators
//--------------------------------------------------------------------------------

template<class S, class PART, class OP>
void pcps::make_time_evolution_op(const pcps::Input & userinp,
                                  const typename pcps::OpBase<S>::const_ptr & op,
                                  std::vector<int> & op_n_sites,
                                  std::vector< boost::shared_array<   int> > & op_sites,
                                  std::vector< boost::shared_array<S> > & op_matrices) {

  // skip the operator if it is of the wrong type
  if (typeid(*op) != typeid(OP)) return;

  // get the number of configurations on a site
  const int configs_per_site = userinp.configs_per_site();

  // create an action value object for the operator
  pcps::action_value_object<S, OP, PART> avo(configs_per_site);

  // check if the operator is contained
  std::vector<int> containing;
  for (int i = 0; i < op_sites.size(); i++) {
    bool contained = true;
    {
      std::set<int> s(op_sites.at(i).get(), op_sites.at(i).get() + op_n_sites.at(i));
      for (int j = 0; j < op->sites().size(); j++)
        contained = ( contained && s.count(op->sites().at(j)) > 0 );
    }
    if (contained) containing.push_back(i);
  }

  // if the operator is not contained in any current time evolution operator, create a new one
  if (containing.size() == 0) {
    const int pos = op_n_sites.size();
    op_n_sites.push_back(op->sites().size());
    op_sites.push_back( pcps::allocate_shared_array<int>(op->sites().size(), 0) );
    for (int i = 0; i < op->sites().size(); i++)
      op_sites[pos][i] = op->sites()[i];
    op_matrices.push_back( pcps::allocate_shared_array<S>(pcps::pow<int>(configs_per_site, 2*op->sites().size()), pcps::zero<S>()) );
    containing.push_back(pos);
  }

  // for each time evolution operator containing the operator, add the operator's contribution
  const S factor = pcps::unity<S>() / double(containing.size());
  for (std::vector<int>::const_iterator c = containing.begin(); c != containing.end(); c++) {

    // get the dimension of the time evolution operator's sites' hilbert space
    const int dim = pcps::pow<int>(configs_per_site, op_n_sites[*c]);

    // get a vector of the time evolution operator's sites
    std::vector<int> tos(op_sites[*c].get(), op_sites[*c].get() + op_n_sites[*c]);

    // get a set of the hamiltonian operator's sites
    std::set<int> hos(op->sites().begin(), op->sites().end());

    // add the operator's contribution to the matrix elements of the time evolution operator
    for (pcps::Cluster bra_occs(userinp.nsites(), configs_per_site, tos); !bra_occs.finished(); bra_occs++)
    for (pcps::Cluster ket_occs(userinp.nsites(), configs_per_site, tos); !ket_occs.finished(); ket_occs++) {

      // check that the bra and ket occupations are compatible on every site the time evolution operator touches
      bool valid = true;
      for (int i = 0; i < op_n_sites[*c]; i++) {

        // get the offset between bra and ket occupation for this site
        int offset = 0;
        if ( hos.count(op_sites[*c][i]) > 0 ) {
          bool found = false;
          for (int j = 0; j < op->sites().size(); j++)
            if (op->sites().at(j) == op_sites[*c][i]) {
              offset = op->offsets().at(j);
              found = true;
            }
          assert( found );
        }

        // check the validity of the site's bra and ket occupations
        valid = ( valid && bra_occs.ket_occ(op_sites[*c][i]) == offset + ket_occs.ket_occ(op_sites[*c][i]) );

      }

      // if the occupations are compatible, add the operators' contribution to the matrix element
      if (valid) {
        int op_occs[4];
        for (int i = 0; i < op->sites().size(); i++)
          op_occs[i] = ket_occs.ket_occ(op->sites().at(i));
        op_matrices[*c][ dim * bra_occs.cmpd() + ket_occs.cmpd() ] += op->coeff() * factor * avo.get_value(op_occs, configs_per_site);
      }

    }

  }

  //std::cout << boost::format("added %20s    to %i time evolution operators") % pcps::type_lookup(typeid(OP)) % containing.size() << std::endl;

}

//--------------------------------------------------------------------------------
// pcps::get_time_evolution_ops -- Convert the hamiltonian's operators into the
//                                 form used by the time evolution routine
//--------------------------------------------------------------------------------

template <class S, class PART>
void pcps::get_time_evolution_ops(const pcps::Input & userinp,
                                  const pcps::Hamiltonian<S> & ham,
                                  std::vector<int> & op_n_sites,
                                  std::vector< boost::shared_array<   int> > & op_sites,
                                  std::vector< boost::shared_array<S> > & op_matrices) {

  // empty the vectors containing the time evolution operators' information
  op_n_sites.clear();
  op_sites.clear();
  op_matrices.clear();

  // create or add to an existing operator for each raising-lowering operator in the hamiltonian
  for (typename std::vector<typename pcps::OpBase<S>::const_ptr>::const_iterator op = ham.ops().begin(); op != ham.ops().end(); op++)
    pcps::make_time_evolution_op<S, PART, pcps::OpRL<S> >(userinp, *op, op_n_sites, op_sites, op_matrices);

  // create or add to an existing operator for each double number operator in the hamiltonian
  for (typename std::vector<typename pcps::OpBase<S>::const_ptr>::const_iterator op = ham.ops().begin(); op != ham.ops().end(); op++)
    pcps::make_time_evolution_op<S, PART, pcps::OpNN<S> >(userinp, *op, op_n_sites, op_sites, op_matrices);

  // create or add to an existing operator for each single number operator in the hamiltonian
  for (typename std::vector<typename pcps::OpBase<S>::const_ptr>::const_iterator op = ham.ops().begin(); op != ham.ops().end(); op++)
    pcps::make_time_evolution_op<S, PART, pcps::OpN<S> >(userinp, *op, op_n_sites, op_sites, op_matrices);

  // create or add to an existing operator for each identity operator in the hamiltonian
  for (typename std::vector<typename pcps::OpBase<S>::const_ptr>::const_iterator op = ham.ops().begin(); op != ham.ops().end(); op++)
    pcps::make_time_evolution_op<S, PART, pcps::OpIdentity<S> >(userinp, *op, op_n_sites, op_sites, op_matrices);

}

//--------------------------------------------------------------------------------
// pcps::get_time_evolution_cluster -- Construct a cluster around a time evolution
//                                     operator.
//--------------------------------------------------------------------------------

template <class S>
void pcps::get_time_evolution_cluster(const pcps::Input & userinp,
                                      const pcps::Wavefunction<S> & wfn,
                                      const int op_n_sites,
                                      const int * const op_sites,
                                      const S * const op_matrices,
                                      boost::shared_ptr<std::vector<pcps::Correlator<S> > > & cluster_corrs,
                                      boost::shared_ptr<std::set<int> > & cluster_sites, 
                                      boost::shared_ptr<std::vector<int> > & cluster_offsets,
                                      int & n_clust_vars,
                                      int & n_clust_configs) {

  // initialize sets for the cluster's sites and correlators
  cluster_sites = boost::shared_ptr<std::set<int> >( new std::set<int>() );
  typename pcps::Correlator<S>::Set cluster_corr_set;

  // add the correlators touched by the operator to the cluster
  for (const int * site = op_sites; site != op_sites + op_n_sites; site++) {
    typename std::vector<pcps::Correlator<S> *>::const_iterator  corr = wfn.touching_correlators(*site).begin();
    typename std::vector<pcps::Correlator<S> *>::const_iterator   end = wfn.touching_correlators(*site).end();
    for ( ; corr != end; corr++) {
      cluster_sites->insert( (**corr).site_vec().begin(), (**corr).site_vec().end() );
      cluster_corr_set.insert(**corr);
    }
  }

  // add all other correlators contained by the cluster's sites
  for (typename std::vector<pcps::Correlator<S> >::const_iterator corr = wfn.correlators().begin(); corr != wfn.correlators().end(); corr++) {
    bool contained = true;
    for (std::vector<int>::const_iterator site = corr->site_vec().begin(); site != corr->site_vec().end(); site++)
      contained = ( contained && cluster_sites->count(*site) > 0 );
    if (contained)
      cluster_corr_set.insert(*corr);
  }

  // create the vector of the cluster's correlators (shallow copies)
  cluster_corrs = boost::shared_ptr<std::vector<pcps::Correlator<S> > >
                    ( new std::vector<pcps::Correlator<S> >(cluster_corr_set.begin(), cluster_corr_set.end()) );

  // determine the offset for each correlator in the cluster
  cluster_offsets = boost::shared_ptr<std::vector<int> >( new std::vector<int>(cluster_corrs->size(), 0) );
  for (int i = 0; i < cluster_corrs->size(); i++) {
    if (cluster_corrs->at(i).fixed()) continue;
    int offset = 0;
    for (typename std::vector<pcps::Correlator<S> >::const_iterator corr = wfn.correlators().begin(); corr != wfn.correlators().end(); corr++) {
      if (corr->fixed()) continue;
      if (*corr == cluster_corrs->at(i))
        cluster_offsets->at(i) = offset;
      offset += corr->nelements();
    }
    assert( offset == wfn.nvars() );
  }

  // determine the number of variables in the cluster's correlators
  n_clust_vars = 0;
  for (typename std::vector<pcps::Correlator<S> >::const_iterator corr = cluster_corrs->begin(); corr != cluster_corrs->end(); corr++)
    if (!corr->fixed())
      n_clust_vars += corr->nelements();

  // determine the number of configurations of the cluster's sites
  n_clust_configs = pcps::pow<int>(userinp.configs_per_site(), cluster_sites->size());

}

//--------------------------------------------------------------------------------
// pcps::time_evolution_driver -- evolves the wavefunction in either real or
//                                imaginary time
//--------------------------------------------------------------------------------

template <class S>
void pcps::time_evolution_driver(const pcps::Input & userinp,
                                 const pcps::Hamiltonian<S> & ham,
                                 pcps::Wavefunction<S> & wfn,
                                 const bool real_time) {

  // start a timer
  pcps::Stopwatch sw; sw.start();

  // get MPI info
  const MPI::Comm & comm = MPI::COMM_WORLD;
  const int nproc = comm.Get_size();
  const int myrank = comm.Get_rank();

  // determine what type of update procudure to use
  const bool join_clusters = userinp.join_trotter_clusters();

  // get the number of sites and the number of configurations on a site
  const int nsites = userinp.nsites();
  const int configs_per_site = userinp.configs_per_site();

  // get timestep and number of iterations
  const double tau = userinp.step_size();
  const int max_iter = userinp.max_macro_loop();
  if (tau <= 0.0 && myrank == 0)
    throw pcps::Exception("timestep must be positive");

  // check that the particle type is acceptable
  if ( myrank == 0 && userinp.particle_type() !=  pcps::keyword::Part_Spin )
    throw pcps::Exception("pcps::time_evolution_driver only implemented for spins");

  // check that the reference type is acceptable
  if ( myrank == 0 && userinp.ref_type() !=  pcps::keyword::Ref_Uniform )
    throw pcps::Exception("pcps::time_evolution_driver only implemented for the uniform reference");

  // if we are performing real time evolution, check that we are using complex numbers
  if (real_time && typeid(S) != typeid(std::complex<double>))
    throw pcps::Exception("real time evolution requires complex numbers");

  // get time evolution operator information
  std::vector<int> op_n_sites;
  std::vector< boost::shared_array<   int> > op_sites;
  std::vector< boost::shared_array<S> > op_matrices;
  pcps::get_time_evolution_ops<S, pcps::Part_Spin>(userinp, ham, op_n_sites, op_sites, op_matrices);
  //if (myrank == 0) {
  //  std::cout << std::endl;
  //  for (int i = 0; i < op_n_sites.size(); i++) {
  //    std::cout << "imag time op on sites ";
  //    for (int j = 0; j < op_n_sites.at(i); j++)
  //      std::cout << boost::format(" %2i") % op_sites[i][j];
  //    std::cout << std::endl;
  //    const int dim = pcps::pow<int>(userinp.configs_per_site(), op_n_sites.at(i));
  //    for (int p = 0; p < dim; p++) {
  //      for (int q = 0; q < dim; q++)
  //        std::cout << boost::format(" %s") % pcps::format_number("%10.2e", op_matrices[i][dim*p+q]);
  //      std::cout << std::endl;
  //    }
  //    std::cout << std::endl;
  //  }
  //  //throw pcps::Exception("stoping here");
  //}

  // multiply update operators by the timestep, and by i for real time evolution
  for (int i = 0; i < op_n_sites.size(); i++)
    pcps::xscal( pcps::pow<int>(configs_per_site, 2*op_n_sites[i]),
                 -(real_time ? pcps::imaginary_unity<S>() : pcps::unity<S>()) * tau,
                 op_matrices[i].get(), 1);

  // determine which of the trotter-decomposed operators belong to this process
  const int nops = op_matrices.size();
  if (nproc > nops)
    throw pcps::Exception( (boost::format("Error in pcps::time_evolution_driver.  More processors than operators. Only have %i operators.") % nops).str() );
  const int extras = nops % nproc;
  const int my_min = myrank * ( nops / nproc ) + std::min(myrank, extras);
  const int my_max = ( myrank + 1 ) * ( nops / nproc ) + std::min(myrank+1, extras);

  // build and store the cluster for each operator
  const int my_nops = my_max - my_min;
  std::vector<boost::shared_ptr<std::vector<pcps::Correlator<S> > > > cluster_corrs(my_nops);
  std::vector<boost::shared_ptr<std::set<int> > > cluster_sites(my_nops);
  std::vector<boost::shared_ptr<std::vector<int> > > cluster_offsets(my_nops);
  std::vector<int> n_clust_vars(my_nops);
  std::vector<int> n_clust_configs(my_nops);

  for (int i = my_min; i < my_max; i++)
    pcps::get_time_evolution_cluster(userinp, wfn, op_n_sites[i], op_sites[i].get(), op_matrices[i].get(),
                                     cluster_corrs[i-my_min], cluster_sites[i-my_min], cluster_offsets[i-my_min],
                                     n_clust_vars[i-my_min], n_clust_configs[i-my_min]);

  // determine the maximum number of variables and configurations in any cluster
  const int max_clust_var    = *std::max_element(n_clust_vars.begin(), n_clust_vars.end());
  const int max_clust_config = *std::max_element(n_clust_configs.begin(), n_clust_configs.end());

  // determine the number of variables in the wavefunction
  const int total_wfn_var = wfn.nvars();

  // create vectors to hold correlators from previous iterations and the difference between previous and current correlators
  std::vector<pcps::Correlator<S> > wfn_prev(wfn.correlators().size());
  std::vector<pcps::Correlator<S> > wfn_diff(wfn.correlators().size());
  for (int i = 0; i < wfn.correlators().size(); i++) {
    wfn_prev.at(i) = pcps::Correlator<S>(wfn.correlators().at(i), true); // deep copy
    wfn_diff.at(i) = pcps::Correlator<S>(wfn.correlators().at(i), true); // deep copy
  }

  // create vectors needed by the time evolution algorithm
  std::vector<S> cluster_tensor_0( max_clust_config                                                              , pcps::zero<S>() );
  std::vector<S> cluster_tensor_1( max_clust_config                                                              , pcps::zero<S>() );
  std::vector<S>                b(                  max_clust_var                                                , pcps::zero<S>() );
  std::vector<S>               AA(    max_clust_var*max_clust_var                                                , pcps::zero<S>() );
  std::vector<S>                A( max_clust_config*max_clust_var                                                , pcps::zero<S>() );
  std::vector<S>                B( max_clust_config*max_clust_var                                                , pcps::zero<S>() );
  std::vector<S>          big_b_0( ( join_clusters ?               total_wfn_var :                           0 ) , pcps::zero<S>() );
  std::vector<S>          big_b_1( ( join_clusters ?               total_wfn_var :                           0 ) , pcps::zero<S>() );
  std::vector<S>         big_AA_0( ( join_clusters ? total_wfn_var*total_wfn_var :                           0 ) , pcps::zero<S>() );
  std::vector<S>         big_AA_1( ( join_clusters ? total_wfn_var*total_wfn_var :                           0 ) , pcps::zero<S>() );
  std::vector<S>                U( ( join_clusters ? total_wfn_var*total_wfn_var : max_clust_var*max_clust_var ) , pcps::zero<S>() );
  std::vector<S>               VT( ( join_clusters ? total_wfn_var*total_wfn_var : max_clust_var*max_clust_var ) , pcps::zero<S>() );
  std::vector<double>              sig( ( join_clusters ?               total_wfn_var :               max_clust_var ) , 0.0                  );

  // create work arrays
  const int lwork = 5 * ( join_clusters ? total_wfn_var : max_clust_var );
  std::vector<S>  work( lwork, pcps::zero<S>() );
  std::vector<double> rwork( lwork, 0.0                  );

  // iterate
  double total_imag_time = 0.0;
  for (int iter = 0; iter < max_iter; iter++) {

    // save the previous iteration's wavefunction and set the differences to zero
    for (int i = 0; i < wfn.correlators().size(); i++) {
      pcps::xcopy(wfn_prev.at(i).nelements(), wfn.correlators().at(i).data_ptr(), 1, wfn_prev.at(i).data_ptr(), 1);
      pcps::xscal(wfn_diff.at(i).nelements(), pcps::zero<S>(), wfn_diff.at(i).data_ptr(), 1);
    }

    // zero the total system AA matrix and b vector
    if (join_clusters) {
      pcps::xscal(total_wfn_var, pcps::zero<S>(), &big_b_0[0], 1);
      pcps::xscal(total_wfn_var, pcps::zero<S>(), &big_b_1[0], 1);
      pcps::xscal(total_wfn_var*total_wfn_var, pcps::zero<S>(), &big_AA_0[0], 1);
      pcps::xscal(total_wfn_var*total_wfn_var, pcps::zero<S>(), &big_AA_1[0], 1);
    }

    // loop over the imaginary time operators in the hamiltonian
    for (int op_num = my_min; op_num < my_max; op_num++) {

      // get the cluster's number
      const int clust_num = op_num - my_min;

      // get references to the cluster's correlators and sites
      std::vector<pcps::Correlator<S> > & corrs = *cluster_corrs[clust_num];
      const std::set<int> & sites = *cluster_sites[clust_num];
      const int nvar = n_clust_vars[clust_num];
      const int nconfig = n_clust_configs[clust_num];
      std::vector<int> & offsets = *cluster_offsets[clust_num];

      // zero the A matrix
      pcps::xscal(nconfig*nvar, pcps::zero<S>(), &A[0], 1);

      // organize the cluster's sites in an order such that the operator's sites appear last
      std::vector<int> site_vec(sites.size());
      {
        std::set<int> op_site_set(op_sites[op_num].get(), op_sites[op_num].get() + op_n_sites[op_num]);
        int i = 0;
        for (std::set<int>::const_iterator site = sites.begin(); site != sites.end(); site++)
          if ( op_site_set.count(*site) == 0 )
            site_vec.at(i++) = *site;
        assert( i + op_n_sites[op_num] == site_vec.size() );
        for (int j = 0; j < op_n_sites[op_num]; j++)
          site_vec.at(i++) = op_sites[op_num][j];
        assert( i == site_vec.size() );
        assert( std::set<int>(site_vec.begin(), site_vec.end()).size() == sites.size() );
      }

      // loop over every cluster configuration to compute the new cluster tensor and A matrix
      for (pcps::Cluster occ(nsites, configs_per_site, site_vec); !occ.finished(); occ++) {

        // get compound index for this configuration
        const int cmpd = occ.cmpd();

        // compute the cluster tensor element
        cluster_tensor_0[cmpd] = pcps::correlator_weight(configs_per_site, corrs, &occ.ket_occs()[0]);

        // compute this configuration's contribution to the A matrix
        int pos = 0;
        for (int i = 0; i < corrs.size(); i++) {
          if (corrs[i].fixed()) continue;
          const int * site = &corrs[i].site_vec()[0];
          const int * site_end = site + corrs[i].site_vec().size();
          int j = 0;
          for ( ; site != site_end; site++)
            j = j * configs_per_site + occ.ket_occ(*site);
          assert( std::abs(corrs[i][j]) > 0.0 );
          A[(pos+j)*nconfig+cmpd] = cluster_tensor_0[cmpd] / corrs[i][j];
          pos += corrs[i].nelements();
        }
        assert( pos == nvar );

      }

      // compute the new cluster configuration via contraction with the time evolution operator
      const int op_dim = pcps::pow<int>(configs_per_site, op_n_sites[op_num]);
      assert( nconfig % op_dim == 0 );
      pcps::xgemm('n', 'n', op_dim, nconfig/op_dim, op_dim, pcps::unity<S>(), &op_matrices[op_num][0], op_dim,
                  &cluster_tensor_0[0], op_dim, pcps::zero<S>(), &cluster_tensor_1[0], op_dim);

      //if (myrank == 0) {
      //  std::cout << std::endl;
      //  std::cout << "cluster sites =" << std::endl;
      //  for (int i = 0; i < site_vec.size(); i++)
      //    std::cout << boost::format(" %2i") % site_vec.at(i);
      //  std::cout << std::endl;
      //  std::cout << std::endl;
      //  std::cout << "new cluster tensor:" << std::endl;
      //  for (int i = 0; i < nconfig; i++)
      //    std::cout << pcps::format_number("%20.6e", cluster_tensor_1[i]) << std::endl;
      //  std::cout << std::endl;
      //}

      // compute AA, which is transpose(A) * A
      pcps::xcopy(nconfig*nvar, &A[0], 1, &B[0], 1);
      pcps::xgemm('c', 'n', nvar, nvar, nconfig, pcps::unity<S>(), &A[0], nconfig, &B[0], nconfig, pcps::zero<S>(), &AA[0], nvar);

      // compute b, which is transpose(A) * (new cluster tensor)
      pcps::xgemm('c', 'n', nvar, 1, nconfig, pcps::unity<S>(), &A[0], nconfig, &cluster_tensor_1[0], nconfig, pcps::zero<S>(), &b[0], nvar);

      // if using the joined cluster update, add this cluster's contribution to the large update matrix and b vector
      if (join_clusters) {

        // add cluster's contribution to full system AA matrix
        int offset_0 = 0;
        for (int i0 = 0; i0 < corrs.size(); i0++) {
          if (corrs[i0].fixed()) continue;
          int offset_1 = 0;
          for (int i1 = 0; i1 < corrs.size(); i1++) {
            if (corrs[i1].fixed()) continue;
            for (int j0 = 0; j0 < corrs[i0].nelements(); j0++)
            for (int j1 = 0; j1 < corrs[i1].nelements(); j1++)
              big_AA_0[ total_wfn_var * ( offsets[i0] + j0 ) + offsets[i1] + j1 ] += AA[ nvar * ( offset_0 + j0 ) + offset_1 + j1 ];
            offset_1 += corrs[i1].nelements();
          }
          assert( offset_1 == nvar );
          offset_0 += corrs[i0].nelements();
        }
        assert( offset_0 == nvar );

        // add cluster's contribution to full system b vector
        offset_0 = 0;
        for (int i0 = 0; i0 < corrs.size(); i0++) {
          if (corrs[i0].fixed()) continue;
          for (int j0 = 0; j0 < corrs[i0].nelements(); j0++)
            big_b_0[ offsets[i0] + j0 ] += b[ offset_0 + j0 ];
          offset_0 += corrs[i0].nelements();
        }
        assert( offset_0 == nvar );

      // otherwise compute and save the update for this cluster's variables
      } else {

        // solve the linear equation
        int info;
        pcps::xgesvd('A', 'A', nvar, nvar, &AA[0], nvar, &sig[0], &U[0], nvar, &VT[0], nvar, &work[0], lwork, &rwork[0], info);
        if (info != 0)
          throw pcps::Exception( (boost::format("xgesvd in pcps::time_evolution_driver failed with error code %i") % info).str() );
        pcps::xgemm('c', 'n', nvar, 1, nvar, pcps::unity<S>(), &U[0], nvar, &b[0], nvar, pcps::zero<S>(), &A[0], nvar);
        for (int i = 0; i < nvar; i++) {
          if (sig[i] > 1.0e-9)
            A[i] /= sig[i];
          else
            A[i] = pcps::zero<S>();
        }
        pcps::xgemm('c', 'n', nvar, 1, nvar, pcps::unity<S>(), &VT[0], nvar, &A[0], nvar, pcps::zero<S>(), &b[0], nvar);

        // update the correlators
        int pos = 0;
        for (int i = 0; i < corrs.size(); i++)
          if (!corrs[i].fixed())
            for (int j = 0; j < corrs[i].nelements(); j++)
              corrs[i][j] += b[pos++];
        assert(pos == nvar);

        // compute and record the difference in each correlator created by this operator, then set the correlator to its previous values
        typename pcps::Correlator<S>::Compare_Obj comp;
        for (int i = 0; i < corrs.size(); i++) {
          pcps::Correlator<S> & cc = corrs[i];
          for (int j = 0; j < wfn_prev.size(); j++) {
            pcps::Correlator<S> & pc = wfn_prev[j];
            pcps::Correlator<S> & dc = wfn_diff[j];
            if (!cc.fixed() && !pc.fixed() && cc == pc) {
              pcps::xaxpy(cc.nelements(), -pcps::unity<S>(), pc.data_ptr(), 1, cc.data_ptr(), 1);
              pcps::xaxpy(cc.nelements(),  pcps::unity<S>(), cc.data_ptr(), 1, dc.data_ptr(), 1);
              pcps::xcopy(cc.nelements(), pc.data_ptr(), 1, cc.data_ptr(), 1);
            }
          }
        }

      }

    }

    // if joining the clusters, solve the total system's least squares equation and update the wavefunction
    if (join_clusters) {

      // add each process's contribution to the total system's AA matrix and b vector together on the root process
      if (typeid(S) == typeid(double)) {
        comm.Reduce(&big_AA_0[0], &big_AA_1[0], total_wfn_var*total_wfn_var, MPI::DOUBLE, MPI::SUM, 0);
        comm.Reduce( &big_b_0[0],  &big_b_1[0],               total_wfn_var, MPI::DOUBLE, MPI::SUM, 0);
      } else if (typeid(S) == typeid(std::complex<double>)) {
        comm.Reduce(&big_AA_0[0], &big_AA_1[0], total_wfn_var*total_wfn_var, MPI::DOUBLE_COMPLEX, MPI::SUM, 0);
        comm.Reduce( &big_b_0[0],  &big_b_1[0],               total_wfn_var, MPI::DOUBLE_COMPLEX, MPI::SUM, 0);
      } else
          throw pcps::Exception("unknown S type in pcps::imag_time_driver");

      // solve the linear equation on the root node and update the wavefunction
      if (myrank == 0) {

        // solve the linear equation
        int info;
        const int n = total_wfn_var;
        pcps::xgesvd('A', 'A', n, n, &big_AA_1[0], n, &sig[0], &U[0], n, &VT[0], n, &work[0], lwork, &rwork[0], info);
        if (info != 0)
          throw pcps::Exception( (boost::format("xgesvd in pcps::time_evolution_driver failed with error code %i") % info).str() );
        pcps::xgemm('c', 'n', n, 1, n, pcps::unity<S>(), &U[0], n, &big_b_1[0], n, pcps::zero<S>(), &big_b_0[0], n);
        for (int i = 0; i < n; i++) {
          if (sig[i] > 1.0e-9)
            big_b_0[i] /= sig[i];
          else
            big_b_0[i] = pcps::zero<S>();
        }
        pcps::xgemm('c', 'n', n, 1, n, pcps::unity<S>(), &VT[0], n, &big_b_0[0], n, pcps::zero<S>(), &big_b_1[0], n);

        // update the wavefunction's correlators
        int pos = 0;
        for (int i = 0; i < wfn.correlators().size(); i++)
          if (!wfn.correlators()[i].fixed())
            for (int j = 0; j < wfn.correlators()[i].nelements(); j++)
              wfn.correlators()[i][j] += big_b_1[pos++];
        assert( pos == wfn.nvars() );

      }

    // if not joining the clusters, combine and apply the separate cluster's updates to the wavefunction
    } else {

      // collect the updates on the root process
      for (int i = 0; i < wfn.correlators().size(); i++) {
        pcps::Correlator<S> & wc = wfn.correlators().at(i);
        pcps::Correlator<S> & dc = wfn_diff.at(i);
        if (wc.fixed()) continue;
        pcps::xscal(wc.nelements(), pcps::zero<S>(), wc.data_ptr(), 1);
        if (typeid(S) == typeid(double))
          comm.Reduce((void *)dc.data_ptr(), (void *)wc.data_ptr(), wc.nelements(), MPI::DOUBLE, MPI::SUM, 0);
        else if (typeid(S) == typeid(std::complex<double>))
          comm.Reduce((void *)dc.data_ptr(), (void *)wc.data_ptr(), wc.nelements(), MPI::DOUBLE_COMPLEX, MPI::SUM, 0);
        else
          throw pcps::Exception("unknown S type in pcps::imag_time_driver");
        pcps::xcopy(wc.nelements(), wc.data_ptr(), 1, dc.data_ptr(), 1);
      }

      // add the updates to the correlators
      if (myrank == 0) {
        for (int i = 0; i < wfn.correlators().size(); i++) {
          pcps::Correlator<S> & wc = wfn.correlators().at(i);
          pcps::Correlator<S> & pc = wfn_prev.at(i);
          pcps::Correlator<S> & dc = wfn_diff.at(i);
          if (wc.fixed()) continue;
          pcps::xcopy(wc.nelements(), pc.data_ptr(), 1, wc.data_ptr(), 1);
          pcps::xaxpy(wc.nelements(), pcps::unity<S>(), dc.data_ptr(), 1, wc.data_ptr(), 1);
          wc.normalize();
        }
      }

    }

    // scale the correlators
    if (myrank == 0 && ( (iter+1) % userinp.dump_freq() == 0 || (iter+1) % userinp.print_freq() == 0 ) )
      wfn.scale_variables();

    // broadcast the new correlators to all processes
    for (typename std::vector<pcps::Correlator<S> >::iterator corr = wfn.correlators().begin(); corr != wfn.correlators().end(); corr++)
      if (!corr->fixed())
        comm.Bcast(corr->data_ptr(), sizeof(S) * corr->nelements(), MPI::CHAR, 0);

    // increment the total imaginary time
    total_imag_time += tau;

    // print info every so many iterations
    if ((iter+1) % userinp.print_freq() == 0) {

      // compute correlation functions and write them to a file
      if (userinp.properties()) {
        const std::string prop_file_name = ( boost::format("property_dump_at_time_%.6f") % total_imag_time ).str();
        pcps::properties<S, pcps::Part_Spin, pcps::Ref_Uniform>(userinp, ham, wfn, userinp.properties_sample_length(), prop_file_name);
      }

      // print energy and correlators
      const S energy = pcps::mc_energy<S, pcps::Part_Spin, pcps::Ref_Uniform>(userinp, ham, wfn);
      if (myrank == 0) {
        std::cout << boost::format("imaginary time       %22.12f        elapsed time = %17.6f sec")
                     % total_imag_time
                     % pcps::get_overall_timer().elapsed_seconds()
                  << std::endl;
        std::cout << boost::format("energy               %22.12f") % pcps::real(energy) << std::endl;
        std::cout << std::endl;
        if (userinp.verbose_print())
          wfn.print(userinp);
      }

    }

    // save wavefunction to disk
    if (myrank == 0 && (iter+1) % userinp.dump_freq() == 0)
      wfn.dump_binary(userinp);

  }

  // save wavefunction to disk
  if (myrank == 0)
    wfn.dump_binary(userinp);

}

// explicitly instantiate the template
template void pcps::time_evolution_driver(const pcps::Input &,
                                          const pcps::Hamiltonian<double> &,
                                          pcps::Wavefunction<double> &,
                                          const bool);
template void pcps::time_evolution_driver(const pcps::Input &,
                                          const pcps::Hamiltonian< std::complex<double> > &,
                                          pcps::Wavefunction< std::complex<double> > &,
                                          const bool);
